#include<stdio.h>
#include<stdlib.h>
#include<FreeImage.h>
#include<cudnn.h>
#include<string.h>
bool SaveImage(const char *fileName, char *buffer, int width, int height)
{
    FREE_IMAGE_FORMAT format = FreeImage_GetFIFFromFilename(fileName);
    FIBITMAP *image = FreeImage_ConvertFromRawBits((BYTE*)buffer, width,
                        height, width * 3, 24,
                        0xFF000000, 0x00FF0000, 0x0000FF00);
    return (FreeImage_Save(format, image, fileName) == TRUE) ? true : false;
}

bool LoadImage(const char*fileName,int &width, int &height,char**data)
{
    FREE_IMAGE_FORMAT format = FreeImage_GetFileType(fileName, 0);
    FIBITMAP* image = FreeImage_Load(format, fileName);

    FIBITMAP* temp = image;
    image = FreeImage_ConvertTo24Bits(image);
    FreeImage_Unload(temp);

    width = FreeImage_GetWidth(image);
    height = FreeImage_GetHeight(image);

    *data = new char[width*height*3];
     memcpy(*data,FreeImage_GetBits(image),width*height*3);
     FreeImage_Unload(image);


}
int main(int argc,char**argv)
{
        char*data;
	int width,height;
	LoadImage("lena.jpeg",width,height,&data);
	printf("width=%d,%d\n",width,height);
	
	float*h_input = new float[width*height*3];
	for(int i =0 ;i<width*height*3;i++){
	   h_input[i] = data[i];
	}
	delete[] data;
   
	cudnnHandle_t cudnn;
	cudnnCreate(&cudnn);

	//input descriptor
	cudnnTensorDescriptor_t input_desc;
	cudnnCreateTensorDescriptor(&input_desc);
	cudnnSetTensor4dDescriptor(input_desc,CUDNN_TENSOR_NHWC,CUDNN_DATA_FLOAT,1,3,width,height);
        //output descriptor
	cudnnTensorDescriptor_t output_desc;
	cudnnCreateTensorDescriptor(&output_desc);
	cudnnSetTensor4dDescriptor(output_desc,CUDNN_TENSOR_NHWC,CUDNN_DATA_FLOAT,1,3,width,height);
	//kernel descriptor
	cudnnFilterDescriptor_t kernel_desc;
	cudnnCreateFilterDescriptor(&kernel_desc);
	cudnnSetFilter4dDescriptor(kernel_desc,CUDNN_DATA_FLOAT,CUDNN_TENSOR_NCHW,3,3,3,3);
       //convolution descriptor
       cudnnConvolutionDescriptor_t convolution_desc;
       cudnnCreateConvolutionDescriptor(&convolution_desc);
       cudnnSetConvolution2dDescriptor(convolution_desc,1,1,1,1,1,1,CUDNN_CROSS_CORRELATION);
	// 
	cudnnConvolutionFwdAlgo_t convolution_algorithm;
	cudnnGetConvolutionForwardAlgorithm(cudnn,input_desc,kernel_desc,convolution_desc,
			output_desc,CUDNN_CONVOLUTION_FWD_PREFER_FASTEST,0,&convolution_algorithm);
	size_t workspace_bytes =0;
	cudnnGetConvolutionForwardWorkspaceSize(cudnn,input_desc,kernel_desc,convolution_desc,output_desc,convolution_algorithm,&workspace_bytes);
	printf("workspace_bytes=%d\n");

	//allocate gpu memory
	void*d_workspace;
	cudaMalloc(&d_workspace,workspace_bytes);

	int imageBytes = 1*3*height*width*sizeof(float);
	float*d_input;
	cudaMalloc(&d_input,imageBytes);
	cudaMemcpy(d_input,h_input,imageBytes,cudaMemcpyHostToDevice);
	float*d_output;
	cudaMalloc(&d_output,imageBytes);
	cudaMemset(d_output,0,imageBytes);

	const float kernel_template[3][3] = {{1,1,1},{1,-8,1},{1,1,1}};
	float h_kernel[3][3][3][3];
	for(int kernel =0;kernel<3;kernel++){
		for(int channel =0;channel<3;channel++){
			for(int row =0;row<3;row++){
				for(int column =0;column<3;column++){
					h_kernel[kernel][channel][row][column] = kernel_template[row][column];
	
				}
			}
		}
	}
	float*d_kernel;
	cudaMalloc(&d_kernel,3*3*3*3*sizeof(float));
	cudaMemcpy(d_kernel,h_kernel,3*3*3*3*sizeof(float),cudaMemcpyHostToDevice);

	const float alpha =1,beta=0;
	cudnnConvolutionForward(cudnn,&alpha,
			input_desc,d_input,
			kernel_desc,d_kernel,
			convolution_desc,convolution_algorithm,
			d_workspace,workspace_bytes,
			&beta,
			output_desc,d_output);

	float *h_output = new float[width*height*3];
	cudaMemcpy(h_output,d_output,imageBytes,cudaMemcpyDeviceToHost);
	char*output = new char[width*height*3];
	for(int i=0;i<width*height*3;i++){
		int temp = h_output[i];
		temp = temp>255?255:temp;
		temp = temp<0?0:temp;
		output[i] = temp;
		
	}
	SaveImage("convolution.jpg",output,width,height);
	delete []h_output;
	delete []output;
	delete []h_input;

	cudaFree(d_kernel);
	cudaFree(d_workspace);
	cudaFree(d_input);
	cudaFree(d_output);
	
	cudnnDestroyTensorDescriptor(input_desc);
	cudnnDestroyTensorDescriptor(output_desc);
	cudnnDestroyFilterDescriptor(kernel_desc);
	cudnnDestroyConvolutionDescriptor(convolution_desc);
	cudnnDestroy(cudnn);

	
}
